import re
from sentence_transformers import SentenceTransformer
import numpy as np

# 加载预训练模型
model = SentenceTransformer("all-MiniLM-L6-v2")


class SemanticChunker:
    def __init__(self, window_size=2, threshold=0.85):
        self.window_size = window_size
        self.threshold = threshold

    def createDocuments(self, text):
        # 分割成句子
        sentences = re.split(r"(。|！|？|\n)", text)
        sents = []
        # 获取分割后的带标题符号的句子
        for i in range(0, len(sentences) - 1, 2):
            s = sentences[i].strip() + sentences[i + 1].strip()
            if s.strip():
                sents.append(s)
        print(f"分割后得到了个{len(sents)}句子")

        docs = []
        # 窗口的开始位置
        start = 0
        while start < len(sents):
            # 计算窗口的结束位置
            end = min(start + self.window_size, len(sents))
            # 获取当前窗口的句子 0:2 0+1 前二行
            window = sents[start:end]
            # 合并窗口里的句子为一个块
            docs.append("".join(window))
            # 移动下一个窗口
            start = end
        print(f"划分为了{len(docs)}块")
        # 计算每个块的向量值
        embeddings = model.encode(docs)
        # 初始化分割点的列表，起始点为0
        split_points = [0]
        # 计算相邻的块之间的相似度
        for i in range(1, len(docs)):  # 1 2
            # 计算余弦相似度
            sim = np.dot(embeddings[i - 1], embeddings[i]) / (
                np.linalg.norm(embeddings[i - 1]) * np.linalg.norm(embeddings[i])
            )
            print(f"sim={sim:.4f}")
            # 如果相似度小于阈值，则作为新的分割点
            if sim < self.threshold:
                split_points.append(i)

        result = []
        for i in range(len(split_points)):
            # 当前 块的起始索引
            start = split_points[i]
            # 当前块的结束索引
            end = split_points[i + 1] if i + 1 < len(split_points) else len(docs)
            # 合并该范围内的窗口为一个块
            chunk = "".join(docs[start:end])
            if chunk.strip():
                result.append(chunk)
        return result


long_text = """ 
今天天气晴朗，适合去公园散步。

量子力学中的叠加态是描述粒子同时处于多个状态的数学工具。

Windows命令行中，复制文件可以使用：copy source.txt destination.txt。

大熊猫主要以竹子为食，是中国的国宝。

欧拉公式被誉为“最美的数学公式”。
"""
semanticChunker = SemanticChunker(window_size=2, threshold=0.85)
documents = semanticChunker.createDocuments(long_text)
for i, doc in enumerate(documents, 1):
    print(f"${i}:{doc}")
    print("-" * 100)
